import torch
import cv2
import numpy as np
from model import MModel


if __name__ == '__main__':
    device = 'cpu'

    model = MModel(input_c=5).to(device)
    model.load_state_dict(torch.load('weights/best.pth', map_location=device))

    dummy = torch.zeros(1, 10, 5)
    torch.onnx.export(
        model, (dummy,), "weights/best.onnx",
        input_names=["input"],
        output_names=["output","feats"],
        dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}, "feats":{0:"batch"}},
        opset_version=11
    )
    print('pth to onnx')
